package geppetto.phraseHMM;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.constrains.ConstrainedProjectionStats;
import geppetto.cat.constrains.SentenceConstrainedProjectionStats;
import geppetto.cat.constrains.SymmetryConstrains;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.SparseTranslationTable;
import geppetto.cat.models.stats.EStepStats;
import geppetto.cat.models.stats.SentenceEStepStats;

import java.io.IOException;




public class SymmetricM1 extends AbstractModel{

	IBMM1 forward;
	IBMM1 backward;
	double optimizationPrecision;
	double slack;
	int _maxStepSize; 
	int _maxNumberIterations;
	
	public SymmetricM1(BilingualCorpus corpus,BilingualCorpus revCorpus, double smoothing,
			double epsilon, double slack,
			int maxStepSize, int maxNumberProjectionIterations) {
		forward = new IBMM1(corpus,smoothing);
		backward = new IBMM1(revCorpus,smoothing);
		_corpus = corpus;
		optimizationPrecision = epsilon;
		this.slack = slack;
		_maxStepSize = maxStepSize;
		_maxNumberIterations = maxNumberProjectionIterations;
	}

	public SymmetricM1(BilingualCorpus corpus,BilingualCorpus revCorpus, 
			SparseTranslationTable tbf,
			SparseTranslationTable tbb,
			double smoothing,
			double epsilon, double slack,
			int maxStepSize, int maxNumberProjectionIterations) {
		forward = new IBMM1(corpus,tbf,smoothing);
		backward = new IBMM1(revCorpus,tbb,smoothing);
		_corpus = corpus;
		optimizationPrecision = epsilon;
		this.slack = slack;
		_maxStepSize = maxStepSize;
		_maxNumberIterations = maxNumberProjectionIterations;
	}
	
	SymmetricM1(IBMM1 forward, IBMM1 backward) {
		this.forward = forward;
		this.backward = backward;
		this._corpus = forward._corpus;
	}
	
	SymmetricM1(){
			
	}
	public String getName() {
		return "Symmetric Constrain M1";
	}

	
	public SentenceConstrainedProjectionStats projectPosteriors(){
		SymmetryConstrains constraints = new SymmetryConstrains(forward, backward, optimizationPrecision, slack,_maxStepSize,_maxNumberIterations);
		SentenceConstrainedProjectionStats stats = constraints.conjugateGradientProjection2();
		//System.out.println(stats.prettyPrint());
		return stats;
	}
	
	public EStepStats createModelStats(){
		ConstrainedProjectionStats pstats = new ConstrainedProjectionStats();
		EStepStats stats = new EStepStats();
		stats.pstats = pstats;
		return stats;
	}
	
	public SentenceEStepStats sentenceEStep(){
		SentenceEStepStats stats = new SentenceEStepStats();
		forward.makeCaches();
		forward.makePosteriors(forward._sSize, forward._fSize);
		backward.makeCaches();
		backward.makePosteriors(backward._sSize, backward._fSize);
		stats.loglikelihood = (forward.getRealLikelihood()+backward.getRealLikelihood())/2;		
		if(stats.loglikelihood < 1.0e-200){
			System.out.println("Skipping a sentence due to likelihood ");
			return stats;
		}
		SentenceConstrainedProjectionStats pstats = projectPosteriors();
		stats.pstats = pstats;
		return stats;
	}
	
	public SymmetricM1 getBackwardModel() {
		SymmetricM1 res = new SymmetricM1(backward, forward);
		res._trained = _trained;
		res._maxNumberIterations = _maxNumberIterations;
		res._maxStepSize = _maxStepSize;
		res.slack = slack;
		res.optimizationPrecision = optimizationPrecision;
		return res;
	}
	
	public EStepStats eStep() {
		forward._currentTrainingIteration = _currentTrainingIteration;
		backward._currentTrainingIteration = _currentTrainingIteration;
		forward.initializeCounts();
		backward.initializeCounts();
		EStepStats stats = createModelStats();
		for (int i = 0; i < forward._nSentences; i++) {
			forward.initSentence(i, BilingualCorpus.TRAIN_CORPUS);
			backward.initSentence(i, BilingualCorpus.TRAIN_CORPUS);
			final int sSize = forward._sourceSentenceIDS.length;
			final int fSize = forward._foreignSentenceIDS.length;
			SentenceEStepStats sstats = sentenceEStep();
			stats.add(sstats,sSize);	
			if(sstats.loglikelihood < 1.0e-200){
				forward._previousSSize = sSize;
				backward._previousSSize = fSize;
				continue;
			}	
			forward.updateFractionalCounts();
			backward.updateFractionalCounts();
			forward._previousSSize = sSize;
			backward._previousSSize = fSize;
		}
		return stats;
	}


	public void finalizeTrain() {
		forward.finalizeTrain();
		backward.finalizeTrain();
	}

	
	public double getNullPhrasePosterior(int sentenceNumber, byte sentenceSource, int[] foreingSentence, int[] sourceSentence, int startForeignIndex, int endForeignIndex) {
		return forward.getNullPhrasePosterior(sentenceNumber, sentenceSource, foreingSentence, sourceSentence, startForeignIndex, endForeignIndex);
	}

	
	public double getPhrasePosterior(int sentenceNumber, byte sentenceSource, int[] foreingSentence, int[] sourceSentence, int startSourceIndex, int endSourceIndex, int startForeignIndex, int endForeignIndex) {
		return forward.getPhrasePosterior(sentenceNumber, sentenceSource, foreingSentence, sourceSentence, startSourceIndex, endSourceIndex, startForeignIndex, endForeignIndex);
	}


	public void initializeTrain() {
		forward.initializeTrain();
		backward.initializeTrain();
		
	}


	public void mStep() {
		forward.mStep();
		backward.mStep();
		
	}
	
	public Alignment posteriorDecodingAlignment(int sentenceNumber, byte sentenceSource, float treshhold, boolean projectPosteriors, ConstrainedProjectionStats stats) {
		forward.initDecoding(sentenceNumber, sentenceSource);
		backward.initDecoding(sentenceNumber, sentenceSource);
		if(projectPosteriors){
			stats.add(projectPosteriors());
		}
		Alignment a = new Alignment(sentenceNumber, sentenceSource, forward._sSize,forward._fSize);
		for (int fIndex = 0; fIndex < forward._fSize; fIndex++) {
			for (int sIndex = 0; sIndex < forward._sSize; sIndex++) {
				if (forward._statePosteriors.getProb(fIndex, sIndex) > treshhold) {
					a.add(sIndex, fIndex);
				}
				a.addPosterior(sIndex, fIndex, forward._statePosteriors.getProb(
						fIndex, sIndex));
			}
		}
		return a;
	}

	public Alignment viterbiAlignment(int sentenceNumber, byte sentenceSource, boolean projectPosteriors, ConstrainedProjectionStats stats) {
		forward.initDecoding(sentenceNumber, sentenceSource);
		backward.initDecoding(sentenceNumber, sentenceSource);
		if(projectPosteriors){
			stats.add(projectPosteriors());
		}
		Alignment a = new Alignment(sentenceNumber, sentenceSource, forward._sSize,forward._fSize);
		for (int foreignPos = 0; foreignPos < forward._fSize; foreignPos++) {
			int pos = 0;
			double max = 0;
			for (int sourcePos = 0; sourcePos < forward._sSize; sourcePos++) {
				double prob = forward._statePosteriors.getProb(foreignPos,sourcePos);
				if (prob > max) {
					max = prob;
					pos = sourcePos;
				}
			}
			double nullProb = forward._statePosteriors.getProb(foreignPos,forward._sSize);
			if (max > nullProb) {
				a.add(pos, foreignPos);
			}
		}
		return a;
		
	}

	public AlignmentsSet viterbiAlignments(byte sentenceSource) {
		return forward.viterbiAlignments(sentenceSource);
	}
	
	public static  Trellis posteriorDifferences(Trellis forward, Trellis backward,int sSize, int fSize){
		Trellis result = new Trellis(fSize,sSize);
		for (int state = 0; state < sSize; state++) {
			for (int pos = 0; pos < fSize; pos++) {
				result.setProb(pos, state, forward.getProb(pos, state)- backward.getProb(state, pos));
			}
		}
		return result;
	}
	
	public static void main(String[] args) throws IOException {
		String corpusDescription = args[0];
		int size = Integer.parseInt(args[1]); // 100k
		int maxSentenceSize = Integer.parseInt(args[2]); // 40
		int numberIterations = Integer.parseInt(args[3]); // 5
		double slack = Double.parseDouble(args[4]); // 5
		double epsilon = Double.parseDouble(args[5]); // 5
		String baseOutput = args[6];
		String outputFile = args[7];
		boolean trainWithResults = Boolean.parseBoolean(args[8]);
		int numberIterationsWithResults = Integer.parseInt(args[9]);
		String saveModelDir = args[10];
		System.out.println("Corpus "+corpusDescription);
		System.out.println("Size " + size);
		System.out.println("Max Sentence size " + maxSentenceSize);
		System.out.println("Number of iterations " + numberIterations);
		System.out.println("Slack " + slack);
		System.out.println("BaseOutpup " + baseOutput);
		System.out.println("OutputFile " + outputFile);
		System.out.println("Train with results " + trainWithResults);
		System.out.println("Number of Iterations with results "
				+ numberIterationsWithResults);

		BilingualCorpus corpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDescription, size, maxSentenceSize);

		BilingualCorpus revCorpus = corpus.reverse();
		
				
		SymmetricM1 m1 = new SymmetricM1(corpus,revCorpus,0,epsilon,slack,1000,1000);
		m1.train(numberIterations,false,saveModelDir);
		SymmetricM1 backm1 = m1.getBackwardModel();
		System.out.println("Done with training");
		
		
		System.out.println("Without decoding");
		
	
		float treshold = m1.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
		AlignmentsSet sa2 = m1.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, treshold,false,false);
		Evaluation eval22 = AlignmentEvaluator.evaluate(sa2, corpus.getGold());
		System.out.println("Forward Posterioir decoding no projection" + eval22);

		
		float tresholdb = backm1.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
		AlignmentsSet sa2b = backm1.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, tresholdb,false,false);
		Evaluation eval22b = AlignmentEvaluator.evaluate(sa2b, revCorpus.getGold());
		System.out.println("Backward  Posterioir decoding no projection" + eval22b);
		
		getCurves(corpus, BilingualCorpus.TEST_CORPUS, false, m1, "fnp");
		getCurves(revCorpus, BilingualCorpus.TEST_CORPUS, false, backm1, "bnp");
		
		System.out.println("With decoding");
		
		
		treshold = m1.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,true);
		sa2 = m1.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, treshold,true,true);
		eval22 = AlignmentEvaluator.evaluate(sa2, corpus.getGold());
		System.out.println("Forward Posterioir decoding with projection" + eval22);

		
		tresholdb = backm1.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,true);
		sa2b = backm1.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, tresholdb,true,true);
		eval22b = AlignmentEvaluator.evaluate(sa2b, revCorpus.getGold());
		System.out.println("Backward  Posterioir decoding with projection" + eval22b);
		
		getCurves(corpus, BilingualCorpus.TEST_CORPUS, true, m1, "fsp");
		getCurves(revCorpus, BilingualCorpus.TEST_CORPUS, true, backm1, "bsp");
		
		
		
	}

	//TODO move to common place
	public static void getCurves(BilingualCorpus corpus, byte sentenceSource, boolean projectPosteriors, AbstractModel model, String grepString){
		int[] sentences = new int[corpus.getNumSentences(sentenceSource)];
		for (int i = 0; i < sentences.length; i++) {
			sentences[i] = i;
		}
		AlignmentsSet predicted = model.posteriorAlignments(sentences,
				sentenceSource, 1f , projectPosteriors,false);
		for(Alignment a : predicted._alignments){
			a.resetPoints();
		}
		int divisions = 100;
		for (int i = 1; i <= divisions; i++) {
			model.predictAux(predicted, 1f * i / divisions);
			float[] results = AlignmentEvaluator.calculateMeasures(predicted,
					corpus.getAlignments(sentences, sentenceSource));
			
			for(Alignment a : predicted._alignments){
				a.resetPoints();
			}
			System.out.println(grepString + " " + 1f * i / divisions + " P: " + results[0] + " R: " + results[1]);
		}
	}
	
	
}
